import os
import SimpleITK as sitk
import pydicom as dicom
from tqdm import tqdm
import numpy as np


def convert_nii_to_dcm(source: str, ref: str, target: str, window=(600, 900), center=(250, 350)):
    files = os.listdir(source)
    for file in tqdm(files):
        source_path = os.path.join(source, file)
        ref_path = os.path.join(ref, file.split('.')[0])
        ref_list = os.listdir(ref_path)
        while len(ref_list) < 10:
            if len(ref_list) == 1:
                ref_path = os.path.join(ref_path, ref_list[0])
            elif 'DICOM' in ref_list:
                ref_path = os.path.join(ref_path, 'DICOM')
            else:
                continue
            ref_list = os.listdir(ref_path)
        target_path = os.path.join(target, file.split('.')[0])
        if os.path.exists(target_path):
            pass
        else:
            os.makedirs(target_path)
        source_data = sitk.GetArrayFromImage(sitk.ReadImage(source_path))
        shape = source_data.shape
        print('   ', file, shape, ref_path)
        if '.' in ref_list[0]:
            ref_list.sort(key=lambda x: int(x.split('.')[4]))
        else:
            ref_list.sort(key=lambda x: int(x))
        slice_min = source_data.min()
        slice_max = source_data.max()
        ref_info_data = dicom.dcmread(os.path.join(ref_path, ref_list[0]))
        final_min = ref_info_data.SmallestImagePixelValue
        final_max = ref_info_data.LargestImagePixelValue
        scale = (final_max - final_min) / (slice_max - slice_min) * 1.0
        for s_i in tqdm(range(shape[0])):
            slice_ = (source_data[shape[0] - s_i - 1] - slice_min) * scale
            name = ref_list[s_i]
            if '.' in name:
                sitk.WriteImage(sitk.GetImageFromArray(slice_.astype(np.uint16)), os.path.join(target_path, name))
                tar = dicom.dcmread(os.path.join(target_path, name))
                os.remove(os.path.join(target_path, name))
            else:
                sitk.WriteImage(sitk.GetImageFromArray(slice_.astype(np.uint16)), os.path.join(target_path, name+'.dcm'))
                tar = dicom.dcmread(os.path.join(target_path, name+'.dcm'))
                os.remove(os.path.join(target_path, name+'.dcm'))
            ref_data = dicom.dcmread(os.path.join(ref_path, name))
            ref_data.PixelData = tar.PixelData
            ref_data.WindowWidth = list(window)
            ref_data.WindowCenter = list(center)
            ref_data.Rows = tar.Rows
            ref_data.Columns = tar.Columns
            ref_data.save_as(os.path.join(target_path, name))


def resample_nii_to_dcm(source: str, ref: str, target: str):
    if os.path.exists(target):
        pass
    else:
        os.makedirs(target)
    files = os.listdir(source)
    for file in tqdm(files):
        source_path = os.path.join(source, file)
        ref_path = os.path.join(ref, file.split('.')[0])
        target_path = os.path.join(target, file)

        source_data = sitk.ReadImage(source_path)
        ref_list = os.listdir(ref_path)
        while len(ref_list) < 10:
            if len(ref_list) == 1:
                ref_path = os.path.join(ref_path, ref_list[0])
            elif 'DICOM' in ref_list:
                ref_path = os.path.join(ref_path, 'DICOM')
            else:
                continue
            ref_list = os.listdir(ref_path)
        xsize, ysize, zsize = source_data.GetSize()
        xspacing, yspacing, zspacing = source_data.GetSpacing()
        origin = source_data.GetOrigin()
        direction = source_data.GetDirection()

        target_depth = len(ref_list)
        target_size = (xsize, ysize, target_depth)
        target_spacing = (xspacing, yspacing, zspacing / (target_depth / float(zsize)))
        target_data = sitk.Resample(source_data, target_size, sitk.Euler3DTransform(), sitk.sitkNearestNeighbor, origin, target_spacing, direction)

        sitk.WriteImage(target_data, target_path)



if __name__ == '__main__':
    source_ = r'xxx\source'
    resample_ = r'xxx\resample'
    ref_ = r'xxx\ref'
    target_ = r'xxx\target'
    resample_nii_to_dcm(source_, ref_, resample_)
    convert_nii_to_dcm(resample_, ref_, target_)
